# definition of the spline functions and prediction routines

cdef extern from "math.h":
    double exp(double)

import numpy as np
cimport numpy as np

# functional terms in the spline function expression
cdef double exp_func(double r, double a1, double a2, double a3) except *:
    return exp(-a1 * r + a2) + a3

cdef double der_exp_func(double r, double a1, double a2) except *:
    return -a1 * exp(-a1 * r + a2)

cdef double cubic_func(double r, double r0, double c0, double c1, double c2, 
                       double c3) except *:
    return c0 + c1 * (r-r0) + c2 * (r-r0)**2  + c3 * (r-r0)**3

cdef double der_cubic_func(double r, double r0, double c1, double c2, 
                           double c3) except *:
    return c1  + 2*c2 * (r-r0)  + 3*c3 * (r-r0)**2

cdef double last_func(double r, double s2, double c0, double c1, double c2,
                      double c3, double c4, double c5) except *:
    return c0 + c1 * (r-s2) + c2 * (r-s2)**2  + c3 * (r-s2)**3 + c4 * (r-s2)**4  + c5 * (r-s2)**5

cdef double der_last_func(double r, double s2, double c1, double c2, double c3,
                          double c4, double c5) except *:
    return c1 + 2*c2 * (r-s2) + 3*c3 * (r-s2)**2 + 4*c4 * (r-s2)**3 + 5*c5 * (r-s2)**4

# base class for complete spline function
cdef class Spline:
    def evaluate(self,
                 double r,
                 double s1, double s2, double cutoff,
                 double a1, double a2, double a3,
                 np.ndarray[np.float64_t, ndim=2] cubic_coeff,
                 np.ndarray[np.float64_t, ndim=2] cubic_intervls,
                 np.ndarray[np.float64_t, ndim=1] last_coeff):
        return 0

cdef class Potential(Spline):
    def evaluate(self,
                 double r,
                 double s1, double s2, double cutoff,
                 double a1, double a2, double a3,
                 np.ndarray[np.float64_t, ndim=2] cubic_coeff,
                 np.ndarray[np.float64_t, ndim=2] cubic_intervls,
                 np.ndarray[np.float64_t, ndim=1] last_coeff):
        if r < s1:
            return exp_func(r, a1, a2, a3)
        elif s2 < r <= cutoff:
            c0 = last_coeff[0]
            c1 = last_coeff[1]
            c2 = last_coeff[2]
            c3 = last_coeff[3]
            c4 = last_coeff[4]
            c5 = last_coeff[5]
            return last_func(r, s2, c0, c1, c2, c3, c4, c5)
        elif r > cutoff:
            return 0.
        else:
            for i, interval in enumerate(cubic_intervls):
                if interval[0] <= r <= interval[1]:
                    r0 = interval[0]
                    c0 = cubic_coeff[i][0]
                    c1 = cubic_coeff[i][1]
                    c2 = cubic_coeff[i][2]
                    c3 = cubic_coeff[i][3]
                    return cubic_func(r, r0, c0, c1, c2, c3)

cdef class Force(Spline):
    def evaluate(self,
                 double r,
                 double s1, double s2, double cutoff,
                 double a1, double a2, double a3,
                 np.ndarray[np.float64_t, ndim=2] cubic_coeff,
                 np.ndarray[np.float64_t, ndim=2] cubic_intervls,
                 np.ndarray[np.float64_t, ndim=1] last_coeff):
        if r < s1:
            return der_exp_func(r, a1, a2)
        elif s2 < r <= cutoff:
            c1 = last_coeff[1]
            c2 = last_coeff[2]
            c3 = last_coeff[3]
            c4 = last_coeff[4]
            c5 = last_coeff[5]
            return der_last_func(r, s2, c1, c2, c3, c4, c5)
        elif r > cutoff:
            return 0.
        else:
            for i, interval in enumerate(cubic_intervls):
                if interval[0] <= r <= interval[1]:
                    r0 = interval[0]
                    c1 = cubic_coeff[i][1]
                    c2 = cubic_coeff[i][2]
                    c3 = cubic_coeff[i][3]
                    return der_cubic_func(r, r0, c1, c2, c3)

                          
# calculation of total force components of system list
def system_2bF_predict(Spline s,
                       bond_dict,
                       B_val,
                       X_val,
                       D_val):
    
    cdef int N_val = len(D_val)
    cdef np.ndarray[np.float64_t, ndim=2] cubic_coeff
    cdef np.ndarray[np.float64_t, ndim=2] cubic_intervls
    cdef np.ndarray[np.float64_t, ndim=1] last_coeff
    cdef int i, iR, L
    cdef np.ndarray[np.float64_t, ndim=1] X_val_i, D_val_i
    cdef Py_ssize_t j, l
    cdef np.ndarray[np.float64_t, ndim=1] result = np.zeros(N_val, 
                                                            dtype=np.float64)
    cdef double temp
    for i in range(N_val):
        iR = int((i-i%3)/3)
        B_val_i = B_val[iR]
        X_val_i = X_val[iR]
        D_val_i = D_val[i]
        L = D_val_i.shape[0]
        temp = 0.
        for l in range(L):
            bond_type = B_val_i[l]
            
            # unpacking spline coeffcients
            cutoff = bond_dict[bond_type]['cutoff']
            a1 = bond_dict[bond_type]['a1']
            a2 = bond_dict[bond_type]['a2']
            a3 = bond_dict[bond_type]['a3']
            cubic_coeff = bond_dict[bond_type]['cubic_coeff']
            cubic_intervls = bond_dict[bond_type]['cubic_intervls']
            s1 = bond_dict[bond_type]['s1']
            s2 = bond_dict[bond_type]['s2']
            last_coeff = bond_dict[bond_type]['last_coeff']
           
            # be sure of the right sign!
            temp += - D_val_i[l] * s.evaluate(X_val_i[l],
                                              s1, s2, cutoff,
                                              a1, a2, a3,
                                              cubic_coeff,
                                              cubic_intervls,
                                              last_coeff)
        result[i] = temp
    
    return result

